"""
func/process_snn.py
利用SNN处理事件数据。
"""
import math
import rospy
import torch
import torch.nn as nn
import os
from typing import Tuple
from rich import print


class Process:
    def __init__(self, device = None) -> None:
        """
        利用脉冲神经网络对运动进行分类。
        @params:
            device: torch.device 计算设备
        """

        # TODO: 完善你的模型和相关参数（开始）

        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 10),
            nn.Sigmoid()
        )
        self.processTimes = [0.0] * 1

        # TODO: 完善你的模型和相关参数（截止）

        self.reset()
    
    def reset(self) -> None:
        """
        重置整个模块。
        """

        # TODO: 如果模型有全局重置，在这里写（开始）

        pass

        # TODO: 如果模型有全局重置，在这里写（截止）

    def time(self) -> float:
        """
        获取当前时间
        @return:
            double 当前的时间戳（单位为秒）
        """
        res = rospy.get_time()
        return res
    
    # TODO: 在此处加入你需要的方法（开始）

    # TODO: 在此处加入你需要的方法（截止）

    def execute(self, data: torch.Tensor) -> torch.Tensor:
        """
        从事件数据中获取避障信息
        @params:
            data: torch.Tensor 事件数据，形状为[n, 5]
        @return:
            torch.Tensor 自定义返回结果
        """

        # TODO: 完成此处的代码（开始）

        print("[green bold]Succeeded![/green bold]")
        print(data.shape)
        res = torch.tensor([0., 0., 0., 0.]).to(data)

        # TODO: 完成此处的代码（截止）

        return res